package com.zhou.framework.aop.aspect;

import com.zhou.framework.redis.util.RedisUtil;
import com.zhou.util.WebUtil;
import com.zhou.framework.aop.anotation.RateLimiter;
import com.zhou.framework.aop.enums.LimitType;
import com.zhou.framework.exception.RateLimiterException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;

/**
 * 接口限流切面
 * @author lang.zhou
 * @since 2023/1/31 17:50
 */
@Aspect
@Slf4j
@Component
public class RateLimiterAspect {

    @Autowired
    private RedisUtil redisUtil;


    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) {
        int time = rateLimiter.time();
        int count = rateLimiter.count();

        String combineKey = getCombineKey(rateLimiter, point);
        try {
            Long number = redisUtil.incr(combineKey, count, time);
            if (number == null || number.intValue() > count){
                throw new RateLimiterException();
            }
        } catch (Exception e) {
            throw new RateLimiterException(e);
        }
    }

    /**
     * 获取限流key
     */
    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
        StringBuilder s = new StringBuilder(rateLimiter.key());
        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if(requestAttributes != null){
            HttpServletRequest request = requestAttributes.getRequest();
            if (rateLimiter.limitType() == LimitType.IP) {
                s.append(WebUtil.getIp(request)).append("-");
            }
        }

        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        s.append(targetClass.getName()).append(".").append(method.getName());
        return s.toString();
    }

}
